

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.SecureRandom;
import java.util.Scanner;

public class AESFileEncryptor {
    // AES算法名称和工作模式
    private static final String AES_ALGORITHM = "AES/CBC/PKCS5Padding";
    // 密钥长度（256位）
    private static final int KEY_SIZE = 256;
    // IV长度（16字节）
    private static final int IV_SIZE = 16;
    
    // ANSI 颜色代码
    private static final String RESET = "\u001B[0m";
    private static final String BLACK = "\u001B[30m";
    private static final String RED = "\u001B[31m";
    private static final String GREEN = "\u001B[32m";
    private static final String YELLOW = "\u001B[33m";
    private static final String BLUE = "\u001B[34m";
    private static final String PURPLE = "\u001B[35m";
    private static final String CYAN = "\u001B[36m";
    private static final String WHITE = "\u001B[37m";
    
    // 主菜单选项
    private static final String[] MENU_OPTIONS = {
        "1. 生成并保存AES密钥",
        "2. 加密文件",
        "3. 解密文件",
        "4. 退出"
    };

    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        
        System.out.println(BLUE + "========================================" + RESET);
        System.out.println(PURPLE + "   欢迎使用 AES-256 文件加密系统     " + RESET);
        System.out.println(BLUE + "========================================" + RESET);
        
        while (true) {
            // 显示主菜单
            displayMenu();
            
            System.out.print(CYAN + "请选择操作 (1-4): " + RESET);
            String choice = scanner.nextLine().trim();
            
            try {
                switch (choice) {
                    case "1":
                        generateAndSaveKey(scanner);
                        break;
                    case "2":
                        encryptFileInteractive(scanner);
                        break;
                    case "3":
                        decryptFileInteractive(scanner);
                        break;
                    case "4":
                        System.out.println(GREEN + "感谢使用 AES-256 文件加密系统！再见！" + RESET);
                        return;
                    default:
                        System.out.println(RED + "无效的选择，请输入 1-4 之间的数字。" + RESET);
                }
            } catch (Exception e) {
                System.out.println(RED + "发生错误: " + e.getMessage() + RESET);
                e.printStackTrace();
            }
            
            System.out.println("\n按回车键继续...");
            scanner.nextLine(); // 等待用户按键继续
            clearConsole();
        }
    }
    
    /**
     * 显示主菜单
     */
    private static void displayMenu() {
        System.out.println(BLUE + "\n============== 主菜单 ==============" + RESET);
        for (String option : MENU_OPTIONS) {
            System.out.println(PURPLE + option + RESET);
        }
        System.out.println(BLUE + "===================================" + RESET);
    }
    
    /**
     * 生成并保存AES密钥
     */
    private static void generateAndSaveKey(Scanner scanner) throws Exception {
        System.out.print(CYAN + "请输入要保存密钥的文件路径 (例如: mykey.bin): " + RESET);
        String keyFilePath = scanner.nextLine().trim();
        
        if (keyFilePath.isEmpty()) {
            System.out.println(RED + "错误: 密钥文件路径不能为空！" + RESET);
            return;
        }
        
        File keyFile = new File(keyFilePath);
        if (keyFile.exists()) {
            System.out.print(YELLOW + "警告: 密钥文件已存在。是否覆盖? (y/n): " + RESET);
            String response = scanner.nextLine().trim().toLowerCase();
            if (!response.equals("y")) {
                System.out.println(GREEN + "操作已取消。" + RESET);
                return;
            }
        }
        
        // 生成随机密钥
        byte[] key = generateKey();
        
        // 保存密钥到文件
        saveKeyToFile(key, keyFilePath);
        
        System.out.println(GREEN + "成功: 密钥已生成并保存到 " + keyFilePath + RESET);
    }
    
    /**
     * 交互式加密文件
     */
    private static void encryptFileInteractive(Scanner scanner) throws Exception {
        System.out.print(CYAN + "请输入密钥文件路径: " + RESET);
        String keyFilePath = scanner.nextLine().trim();
        
        System.out.print(CYAN + "请输入要加密的文件路径: " + RESET);
        String inputFile = scanner.nextLine().trim();
        
        System.out.print(CYAN + "请输入加密结果保存路径: " + RESET);
        String outputFile = scanner.nextLine().trim();
        
        // 验证输入
        if (keyFilePath.isEmpty() || inputFile.isEmpty() || outputFile.isEmpty()) {
            System.out.println(RED + "错误: 所有路径都不能为空！" + RESET);
            return;
        }
        
        // 加载密钥
        byte[] key = loadKeyFromFile(keyFilePath);
        
        // 加密文件
        encryptFile(key, inputFile, outputFile);
        
        System.out.println(GREEN + "成功: 文件加密完成。加密结果已保存到 " + outputFile + RESET);
    }
    
    /**
     * 交互式解密文件
     */
    private static void decryptFileInteractive(Scanner scanner) throws Exception {
        System.out.print(CYAN + "请输入密钥文件路径: " + RESET);
        String keyFilePath = scanner.nextLine().trim();
        
        System.out.print(CYAN + "请输入要解密的加密文件路径: " + RESET);
        String inputFile = scanner.nextLine().trim();
        
        System.out.print(CYAN + "请输入解密结果保存路径: " + RESET);
        String outputFile = scanner.nextLine().trim();
        
        // 验证输入
        if (keyFilePath.isEmpty() || inputFile.isEmpty() || outputFile.isEmpty()) {
            System.out.println(RED + "错误: 所有路径都不能为空！" + RESET);
            return;
        }
        
        // 加载密钥
        byte[] key = loadKeyFromFile(keyFilePath);
        
        // 解密文件
        decryptFile(key, inputFile, outputFile);
        
        System.out.println(GREEN + "成功: 文件解密完成。解密结果已保存到 " + outputFile + RESET);
    }

    /**
     * 生成256位的AES随机密钥
     */
    private static byte[] generateKey() throws Exception {
        SecureRandom random = new SecureRandom();
        byte[] key = new byte[KEY_SIZE / 8]; // 256位=32字节
        random.nextBytes(key);
        return key;
    }

    /**
     * 将密钥保存到文件
     */
    private static void saveKeyToFile(byte[] key, String filePath) throws IOException {
        try (FileOutputStream fos = new FileOutputStream(filePath)) {
            fos.write(key);
        }
    }

    /**
     * 从文件加载AES密钥
     */
    private static byte[] loadKeyFromFile(String filePath) throws IOException {
        return Files.readAllBytes(Paths.get(filePath));
    }

    /**
     * 使用AES加密文件
     */
    private static void encryptFile(byte[] key, String inputFile, String outputFile) throws Exception {
        // 读取原始文件内容
        byte[] fileContent = Files.readAllBytes(Paths.get(inputFile));
        
        // 创建CBC模式使用的IV参数
        byte[] iv = new byte[IV_SIZE];
        new SecureRandom().nextBytes(iv);
        
        // 创建AES密钥规格
        SecretKeySpec keySpec = new SecretKeySpec(key, "AES");
        IvParameterSpec ivSpec = new IvParameterSpec(iv);
        
        // 创建并初始化加密器
        Cipher cipher = Cipher.getInstance(AES_ALGORITHM);
        cipher.init(Cipher.ENCRYPT_MODE, keySpec, ivSpec);
        
        // 执行加密
        byte[] encrypted = cipher.doFinal(fileContent);
        
        // 将IV和加密后的数据合并，写入输出文件
        try (FileOutputStream fos = new FileOutputStream(outputFile)) {
            fos.write(iv); // 先写入IV
            fos.write(encrypted); // 再写入加密数据
        }
    }

    /**
     * 使用AES解密文件
     */
    private static void decryptFile(byte[] key, String inputFile, String outputFile) throws Exception {
        // 读取加密文件内容
        byte[] encryptedData = Files.readAllBytes(Paths.get(inputFile));
        
        // 提取IV（前16字节）
        byte[] iv = new byte[IV_SIZE];
        System.arraycopy(encryptedData, 0, iv, 0, IV_SIZE);
        IvParameterSpec ivSpec = new IvParameterSpec(iv);
        
        // 提取实际加密数据（剩余部分）
        int encryptedFileSize = encryptedData.length - IV_SIZE;
        byte[] encryptedContent = new byte[encryptedFileSize];
        System.arraycopy(encryptedData, IV_SIZE, encryptedContent, 0, encryptedFileSize);
        
        // 创建AES密钥规格
        SecretKeySpec keySpec = new SecretKeySpec(key, "AES");
        
        // 创建并初始化解密器
        Cipher cipher = Cipher.getInstance(AES_ALGORITHM);
        cipher.init(Cipher.DECRYPT_MODE, keySpec, ivSpec);
        
        // 执行解密
        byte[] decrypted = cipher.doFinal(encryptedContent);
        
        // 写入解密后的内容到输出文件
        try (FileOutputStream fos = new FileOutputStream(outputFile)) {
            fos.write(decrypted);
        }
    }
    
    /**
     * 清除控制台屏幕
     */
    private static void clearConsole() {
        // 对于Windows系统
        if (System.getProperty("os.name").contains("Windows")) {
            try {
                new ProcessBuilder("cmd", "/c", "cls").inheritIO().start().waitFor();
            } catch (Exception e) {
                // 忽略异常
            }
        } else {
            // 对于Unix/Linux/Mac系统
            System.out.print("\033[H\033[2J");
            System.out.flush();
        }
    }
}